#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/extensions/XShm.h>
#include <sys/shm.h>

#include <vector>
#include <thread>
#include <queue>
#include <mutex>
#include <cstring>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>

#include "tcpserver.hpp"
#include "semaphore.hpp"
#include "imageProcess.hpp"

#define Height 800
#define Width 1280
#define SegmentWidth (64)
#define SegmentHeight (40)

int PORT = 6118;
int Quality = 80;

static bool is_running = True;
static Display *display;
static Window root;
static XImage *ximage;

static Semaphore sem_grab;
static Semaphore sem_send;

static std::mutex mtx_grab;
static std::mutex mtx_compress;

#pragma pack(push, 8) // 设置结构体按8字节对齐
static uchar frame_buffer[800 * 1280 * 4];
static uchar frame_buffer_copy[800 * 1280 * 4];
#pragma pack(pop) // 恢复默认对齐方式

static std::vector<ImageBlock> image_blocks;
static std::vector<uchar> jpg_image;
static std::vector<uchar> data_stream;

int handleFunction(int client_fd, std::string cmd)
{
    // cv::Mat gray;
    // std::vector<ImageBlock> image_blocks_copy;
    std::vector<uchar> _data_stream;

    uint32_t left;
    uint32_t top;
    uint32_t width;
    uint32_t height;
    uint32_t jpeg_len;

    // 创建一个全黑的图像矩阵
    mtx_grab.lock();
    memset(frame_buffer_copy, 0, sizeof(frame_buffer_copy));
    mtx_grab.unlock();

    try
    {
        Quality = std::stoi(cmd);

        while (is_running)
        {
            auto start = std::chrono::high_resolution_clock::now();

            sem_send.wait_for_signal();

            mtx_compress.lock();
            std::vector<uchar> _data_stream = data_stream;
            mtx_compress.unlock();

            // // mtx_grab.lock();
            // // memcpy(frame_buffer_copy, frame_buffer, sizeof(frame_buffer));
            // // cv::Mat image(Height, Width, CV_8UC4, (char *)frame_buffer_copy);
            // // mtx_grab.unlock();

            // auto start = std::chrono::high_resolution_clock::now();

            // mtx_grab.lock();
            // // memcopy(frame_buffer_copy, frame_buffer, Height * Width * sizeof(uint32_t));
            // imageBlockingCompareAndCopy(frame_buffer, frame_buffer_copy, Width, Height, SegmentWidth, SegmentHeight, image_blocks);
            // cv::Mat image(Height, Width, CV_8UC4, (char *)frame_buffer_copy);
            // // image_blocks_copy = image_blocks;
            // mtx_grab.unlock();

            // auto end = std::chrono::high_resolution_clock::now();
            // std::chrono::duration<double, std::milli> elapsed = end - start;
            // std::cout << "preprocess image Elapsed time: " << elapsed.count() << " ms\n";

            // sem_grab.send_signal();
            // /***********************************************************************/
            // auto _start = std::chrono::high_resolution_clock::now();
            // data_stream.clear();
            // for (std::vector<ImageBlock>::iterator it = image_blocks.begin(); it != image_blocks.end(); ++it)
            // {
            //     // 定义矩形区域 (x, y, Width, height) 截取图像
            //     cv::Rect rect(it->left, it->top, it->width, it->height);
            //     cv::Mat roi = image(rect);
            //     // 编码图像为 JPEG 字节流
            //     std::vector<uchar> _jpeg_buffer;
            //     _jpeg_buffer.clear();
            //     cv::imencode(".jpg", roi, _jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));
            //     // std::cout << "jpg " << it->left << " " << it->top << " data len: " << _jpeg_buffer.size() << " \n";
            //     jpeg_len = _jpeg_buffer.size();
            //     data_stream.insert(data_stream.end(), reinterpret_cast<uchar *>(&*it), reinterpret_cast<uchar *>(&*it) + sizeof(ImageBlock));
            //     data_stream.insert(data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
            //     data_stream.insert(data_stream.end(), _jpeg_buffer.begin(), _jpeg_buffer.end());
            // }

            // auto _end = std::chrono::high_resolution_clock::now();
            // std::chrono::duration<double, std::milli> _elapsed = _end - _start;
            // std::cout << "imencode block jpeg Elapsed time: " << _elapsed.count() << " ms\n";

            /***********************************************************************/
            // auto start = std::chrono::high_resolution_clock::now();
            // // // // 定义矩形区域 (x, y, Width, height) 截取图像
            // // // cv::Rect roi(60, 120, 720, 500);
            // // // cv::Mat cropped_image = image(roi);

            // // 编码图像为 JPEG 字节流
            // std::vector<uchar> jpeg_buffer;
            // cv::imencode(".jpg", image, jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));

            // auto end = std::chrono::high_resolution_clock::now();
            // std::chrono::duration<double, std::milli> elapsed = end - start;
            // std::cout << "imencode whole jpeg Elapsed time: " << elapsed.count() << " ms\n";

            // jpeg_len = jpeg_buffer.size();
            // data_stream.clear();
            // ImageBlock ib = {0, 0, Width, Height};
            // data_stream.insert(data_stream.end(), reinterpret_cast<uchar *>(&ib), reinterpret_cast<uchar *>(&ib) + sizeof(ImageBlock));
            // data_stream.insert(data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
            // data_stream.insert(data_stream.end(), jpeg_buffer.begin(), jpeg_buffer.end());

            // uint32_t len = jpeg_buffer.size();
            // if (send(client_fd, (uchar *)&len, sizeof(uint32_t), MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
            //     return 0;

            // if (send(client_fd, jpeg_buffer.data(), len, MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
            //     return 0;

            /***********************************************************************/
            uint32_t len = _data_stream.size();
            if (send(client_fd, (uchar *)&len, sizeof(uint32_t), MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
                return 0;

            if (len == 0)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
                continue;
            }

            if (send(client_fd, _data_stream.data(), _data_stream.size(), MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
                return 0;

            auto end = std::chrono::high_resolution_clock::now();
            std::chrono::duration<double, std::milli> elapsed = end - start;
            std::cout << "preprocess image Elapsed time: " << elapsed.count() << " ms\n";

            std::cout << "\n";
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Exception: " << e.what() << std::endl;
        return 0;
    }

    // close(client_fd);
    return 1;
}

void compressThread()
{
    uint32_t left;
    uint32_t top;
    uint32_t width;
    uint32_t height;
    uint32_t jpeg_len;
    std::vector<uchar> _data_stream;

    memset(frame_buffer_copy, 0xff, sizeof(frame_buffer_copy));

    std::this_thread::sleep_for(std::chrono::seconds(1));
    while (is_running)
    {
        auto start = std::chrono::high_resolution_clock::now();

        mtx_grab.lock();
        // memcopy(frame_buffer_copy, frame_buffer, Height * Width * sizeof(uint32_t));
        imageBlockingCompareAndCopy(frame_buffer, frame_buffer_copy, Width, Height, SegmentWidth, SegmentHeight, image_blocks);
        cv::Mat image(Height, Width, CV_8UC4, (char *)frame_buffer_copy);
        mtx_grab.unlock();

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "preprocess image Elapsed time: " << elapsed.count() << " ms\n";

        sem_grab.send_signal();
        auto _start = std::chrono::high_resolution_clock::now();

        /***********************************************************************/
        _data_stream.clear();
        for (std::vector<ImageBlock>::iterator it = image_blocks.begin(); it != image_blocks.end(); ++it)
        {
            // 定义矩形区域 (x, y, Width, height) 截取图像
            cv::Rect rect(it->left, it->top, it->width, it->height);
            cv::Mat roi = image(rect);
            // 编码图像为 JPEG 字节流
            std::vector<uchar> _jpeg_buffer;
            _jpeg_buffer.clear();
            cv::imencode(".jpg", roi, _jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));
            // std::cout << "jpg " << it->left << " " << it->top << " data len: " << _jpeg_buffer.size() << " \n";
            jpeg_len = _jpeg_buffer.size();
            _data_stream.insert(_data_stream.end(), reinterpret_cast<uchar *>(&*it), reinterpret_cast<uchar *>(&*it) + sizeof(ImageBlock));
            _data_stream.insert(_data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
            _data_stream.insert(_data_stream.end(), _jpeg_buffer.begin(), _jpeg_buffer.end());
        }
        if (image_blocks.size() == 0)
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(25));
        }

        // /**********************************************/
        // // // 编码图像为 JPEG 字节流
        // // std::vector<uchar> jpeg_buffer;
        // // cv::imencode(".jpg", image, jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));

        // // jpeg_len = jpeg_buffer.size();
        // // _data_stream.clear();
        // // ImageBlock ib = {0, 0, Width, Height};
        // // _data_stream.insert(_data_stream.end(), reinterpret_cast<uchar *>(&ib), reinterpret_cast<uchar *>(&ib) + sizeof(ImageBlock));
        // // _data_stream.insert(_data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
        // // _data_stream.insert(_data_stream.end(), jpeg_buffer.begin(), jpeg_buffer.end());
        // /**********************************************/

        auto _end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> _elapsed = _end - _start;
        std::cout << "imencode block jpeg Elapsed time: " << _elapsed.count() << " ms\n";

        mtx_compress.lock();
        data_stream = _data_stream;
        mtx_compress.unlock();

        // sem_grab.send_signal();
        sem_send.send_signal();
    }
}

int main()
{
    TCPServer server(PORT, handleFunction);
    server.start();
    std::cout << "Server is running on port " << PORT << "..." << std::endl;

    XShmSegmentInfo shminfo;
    display = XOpenDisplay(NULL);
    root = DefaultRootWindow(display);

    XWindowAttributes attributes;
    XGetWindowAttributes(display, root, &attributes);

    // 创建共享内存图像
    ximage = XShmCreateImage(display,
                             DefaultVisual(display, DefaultScreen(display)),
                             attributes.depth,
                             ZPixmap,
                             NULL,
                             &shminfo,
                             attributes.width, attributes.height);

    // XShmSegmentInfo shminfo;
    shminfo.shmid = shmget(IPC_PRIVATE, ximage->bytes_per_line * ximage->height, IPC_CREAT | 0777);
    shminfo.shmaddr = ximage->data = (char *)shmat(shminfo.shmid, 0, 0);
    shminfo.readOnly = False;

    XShmAttach(display, &shminfo);

    // // 设置图像的宽度和高度
    // int width = 640;
    // int height = 480;

    // // 创建一个全黑的图像矩阵
    // cv::Mat frame(height, width, CV_8UC3, cv::Scalar(0, 0, 0));

    std::thread compress_thread(compressThread);

    while (is_running)
    {
        sem_grab.wait_for_signal();

        XSync(display, False);
        XShmGetImage(display, root, ximage, 0, 0, AllPlanes);

        if (Height != ximage->height || Width != ximage->width)
        {
            std::cout << "Screen resolution error: " << ximage->width << "x" << ximage->height << "\n";
            continue;
        }

        auto start = std::chrono::high_resolution_clock::now();

        mtx_grab.lock();
        memcopy(frame_buffer, (uchar *)ximage->data, ximage->bytes_per_line * ximage->height);
        // imageBlockingCompareAndCopy((uchar *)ximage->data, frame_buffer, Width, Height, SegmentWidth, SegmentHeight, image_blocks);

        // int byte_len = ximage->bytes_per_line * ximage->height;
        // memcpy(frame_buffer, (char *)ximage->data, byte_len);
        mtx_grab.unlock();
        // sem_grab.send_signal();

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "XShmGetImage Elapsed time: " << elapsed.count() << " ms\n";

        // std::vector<uchar> frame_buffer_copy;
        // frame_buffer_copy.clear();
        // frame_buffer_copy.reserve(frame_buffer.size()); // 预分配空间
        // std::copy(frame_buffer.begin(), frame_buffer.end(), std::back_inserter(frame_buffer_copy));
        // cv::Mat image(height, Width, CV_8UC4, (char *)frame_buffer_copy.data());

        // // 编码图像为 JPEG 字节流
        // std::vector<uchar> jpeg_buffer;
        // cv::imencode(".jpg", image, jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));

        // cv::Mat image(ximage->height, ximage->width, CV_8UC4, (char *)ximage->data, ximage->bytes_per_line);

        // // 编码图像为 JPEG 字节流
        // std::vector<uchar> jpeg_buffer;
        // cv::imencode(".jpg", image, jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));

        // auto end = std::chrono::high_resolution_clock::now();
        // std::chrono::duration<double, std::milli> elapsed = end - start;
        // std::cout << "XShmGetImage Elapsed time: " << elapsed.count() << " ms\n";
    }

    shmdt(shminfo.shmaddr);
    XShmDetach(display, &shminfo);
    XDestroyImage(ximage);
    XCloseDisplay(display);

    server.stop();
    std::cout << "Server stopped." << std::endl;

    return 0;
}

int main_legacy()
{
    Display *display = XOpenDisplay(NULL);
    Window root = DefaultRootWindow(display);

    XWindowAttributes attributes;
    XGetWindowAttributes(display, root, &attributes);

    XImage *image = XGetImage(display, root, 0, 0, attributes.width, attributes.height, AllPlanes, ZPixmap);

    int width = image->width;
    int height = image->height;

    std::cout << "Captured Image Size: " << Width << "x" << height << std::endl;

    // 在这里可以对image进行进一步处理，如保存为文件或进行图像操作

    XDestroyImage(image);
    XCloseDisplay(display);

    return 0;
}